You write custom CUDA kernels to replace PyTorch operators for speedups.
Implement Group-Softmax-Affine gating on NCHW: Given x[N,C,H,W], per-channel scale[C], bias[C], and an integer group size g, compute z = x*scale[c] + bias[c], then for each (n,h,w) apply softmax within channel groups of size g: g_k = softmax(z over channels in group k), and output y = g_k * x for those channels. The CUDA kernel should process one spatial location per block, cache x and z in shared memory, compute per-group max/sub-exp/sum, and write gated outputs in one pass. Provide a PyTorch reference with nn.Parameter scale, bias, and group as Parameter. Accuracy must match within rtol=1e-3.
